#!/bin/bash

# ============================================================
# NVIDIA GPU Operator & Monitoring 自动化部署脚本 (简化引用版)
# ============================================================

# 出错立即停止
set -e

# 颜色定义
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m' # No Color

# 变量定义
OFFLINE_IMAGE_FILE="../offline_app_images/gpu-operator.v22.9.2.tar.gz"

# 标准短名称 (Sealos 会自动匹配 localhost/ 或 docker.io/ 前缀)
TARGET_IMAGE_NAME="labring/gpu-operator:v22.9.2"

# 关键词定义 (用于搜索)
# Step 1 搜索关键词
IMAGE_SEARCH_KEY="gpu-operator"
IMAGE_SEARCH_TAG="v22.9.2"

# Step 2 旧插件关键词
OLD_PLUGIN_KEYWORD="nvidia-k8s-device-plugin"

# Step 3 GPU Operator 安装状态关键词
OPERATOR_INSTALLED_KEYWORD="gpu-operator"

# 资源文件
SM_YAML="./nvidia-gpu-servicemonitor.yml"
RULE_YAML="./prometheus-gpu-usage.yml"
DASHBOARD_JSON="./dcgm-exporter-dashboard.json"
GRAFANA_NS="cattle-dashboards"

# 辅助函数
log() { echo -e "${GREEN}[$(date +'%T')] $1${NC}"; }
warn() { echo -e "${YELLOW}[$(date +'%T')] $1${NC}"; }
error() { echo -e "${RED}[$(date +'%T')] $1${NC}"; }

check_file() {
    if [[ ! -f "$1" ]]; then
        error "错误：文件不存在 -> $1"
        exit 1
    fi
}

# 获取当前 sealos 已安装列表 (去除首尾空格)
get_installed_images() {
    sealos list --format " {{.ImageName}}" | awk '{$1=$1;print}'
}

# ============================================================
# 1. 加载映像 (不再纠结完整名称，只检查是否存在)
# ============================================================
log "1. 检查本地映像库状态..."

# 只要 grep 能匹配到 gpu-operator 和 v22.9.2 即可，忽略前缀
if sealos images | grep "${IMAGE_SEARCH_KEY}" | grep -q "${IMAGE_SEARCH_TAG}"; then
    log "  -> 检测到本地已存在相关映像，跳过加载。"
else
    log "  -> 本地未找到映像，准备加载离线包..."
    check_file "$OFFLINE_IMAGE_FILE"
    sealos load -i "$OFFLINE_IMAGE_FILE"
    log "  -> 映像加载完成。"
fi

# ============================================================
# 2. 卸载旧版插件 (通过 list 获取准确名称)
# ============================================================
log "2. 检查是否安装了旧版 nvidia-k8s-device-plugin..."

INSTALLED_LIST=$(get_installed_images)
# 查找旧插件的准确安装名称 (如 localhost/labring/nvidia-k8s-device-plugin:v0.15.0)
OLD_PLUGIN_FULL_NAME=$(echo "$INSTALLED_LIST" | grep "${OLD_PLUGIN_KEYWORD}" | head -n 1 || true)

if [[ -n "$OLD_PLUGIN_FULL_NAME" ]]; then
    warn "  -> 检测到旧版插件已安装: ${OLD_PLUGIN_FULL_NAME}"
    warn "  -> 正在执行卸载..."
    
    # 卸载时使用 list 中查到的原名，以确保 sealos 能正确识别要卸载的实例
    sealos run "${OLD_PLUGIN_FULL_NAME}" -e uninstall=true
    log "  -> 卸载命令已执行。"
else
    log "  -> 未检测到旧版插件安装记录，跳过卸载。"
fi

# ============================================================
# 3. 安装 GPU Operator (通过 list 检查状态)
# ============================================================
log "3. 检查 GPU Operator 安装状态..."

# 刷新已安装列表
INSTALLED_LIST=$(get_installed_images)

if echo "$INSTALLED_LIST" | grep -q "${OPERATOR_INSTALLED_KEYWORD}"; then
    warn "  -> 检测到 GPU Operator 已在安装列表中，跳过安装步骤。"
else
    log "  -> 未检测到 GPU Operator，开始安装..."
    log "  -> 使用标准名称: ${TARGET_IMAGE_NAME}"
    
    # 直接使用短名称安装，Sealos 会自动匹配 localhost/docker.io
    sealos run "${TARGET_IMAGE_NAME}" -e HELM_OPTS="--set toolkit.enabled=false"
    
    # 融合等待步骤
    warn "  -> (融合步骤) 等待 120秒，让 Operator 初始化并拉起组件..."
    sleep 120
fi

# ============================================================
# 4. 配置 ServiceMonitor
# ============================================================
log "4. 配置 NVIDIA GPU ServiceMonitor..."
check_file "$SM_YAML"
kubectl apply -f "$SM_YAML"
sleep 5

# ============================================================
# 5. 配置 Prometheus 规则
# ============================================================
log "5. 配置 Prometheus GPU 占用率统计规则..."
check_file "$RULE_YAML"
kubectl apply -f "$RULE_YAML"

warn "  -> 等待 60秒，让 Prometheus 重新加载配置并发现 Target..."
sleep 60

# ============================================================
# 6. 配置 Grafana Dashboard
# ============================================================
log "6. 创建并加载 Grafana Dashboard..."
check_file "$DASHBOARD_JSON"

# 清理旧 ConfigMap
kubectl delete configmap dcgm-exporter-dashboard -n "$GRAFANA_NS" --ignore-not-found

# 创建 ConfigMap
kubectl create configmap dcgm-exporter-dashboard \
  --from-file="$DASHBOARD_JSON" \
  -n "$GRAFANA_NS"

# 打标签
kubectl label configmap dcgm-exporter-dashboard \
  grafana_dashboard=1 \
  -n "$GRAFANA_NS" --overwrite

## 自动获取 Grafana 的 ServiceAccount 名称
#SA_NAME=$(kubectl get sa -n "$GRAFANA_NS" -l app.kubernetes.io/name=grafana -o jsonpath="{.items[0].metadata.name}")
#
## 创建 RoleBinding (赋予读取权限)
## 这里使用 cluster-admin 仅限于该 namespace，确保权限绝对够用
#kubectl create rolebinding grafana-dashboard-access \
#  --clusterrole=cluster-admin \
#  --serviceaccount="${GRAFANA_NS}:${SA_NAME}" \
#  -n "$TARGET_NS" \
#  --dry-run=client -o yaml | kubectl apply -f -

log "============================================"
log "✅ 部署完成！"
log "请参照文档进入  Web 界面验证 Dashboard。"
log "============================================"
